# !/usr/bin/env python3
# coding=utf8
"""
"""
import argparse
import datetime
import enum
import logging
import peewee
import playhouse.sqlite_ext
from typing import Any, Dict, List, Set, Tuple, Type, Optional, Union, Callable
import vnpy.trader.constant
import vnpy.trader.database
import vnpy.tools.database_cta
import vnpy.trader.object
import vnpy.trader.utility
import vnpy_sqlite
from vnpy.tools.chinese_mainland_time_config import ChineseMainlandTimeConfig


def vnpy_sqlite_reinitialize(db0: peewee.SqliteDatabase, sqlite_filename: str):
    """"""
    if sqlite_filename:
        db1 = vnpy_sqlite.sqlite_database.db
        db2 = vnpy_sqlite.sqlite_database.DbBarData._meta.database
        assert isinstance(db0, peewee.SqliteDatabase)
        assert isinstance(db1, peewee.SqliteDatabase)
        assert isinstance(db2, peewee.SqliteDatabase)
        assert id(db0) == id(db1) == id(db2)
        #
        path: str = str(vnpy.trader.utility.get_file_path(filename_dst))
        db0.init(database=path)
        #
        db0.connect()
        db0.create_tables(
            [
                vnpy_sqlite.sqlite_database.DbBarData,
                vnpy_sqlite.sqlite_database.DbTickData,
                vnpy_sqlite.sqlite_database.DbBarOverview,
                vnpy_sqlite.sqlite_database. DbTickOverview,
            ]
        )


class TdxPeriod(enum.Enum):
    """
    通达信的Period
    """
    _00_1min = "1min"
    _01_5min = "5min"
    _02_15min = "15min"
    _03_30min = "30min"
    _04_60min = "60min"
    _05_day = "day"
    _06_week = "week"
    _07_month = "month"
    _08_multi_min = "multi-min"
    _09_multi_day = "multi-day"
    _10_quarter = "quarter"
    _11_year = "year"
    _12_5sec = "5sec"
    _13_multi_sec = "multi-sec"


_tdx_path: str = str(vnpy.trader.utility.get_file_path("libTdxDataSaver_zxcvbnm.db"))
_tdx_db: playhouse.sqlite_ext.SqliteExtDatabase = playhouse.sqlite_ext.SqliteExtDatabase(_tdx_path)


class DbBarDataTdx(peewee.Model):
    """通过lib将Tdx的kBar数据download到SQLite的数据表"""
    # when using the default parameter, the values are set by Peewee rather than being a part of the actual table and column definition.
    # 所以使用peewee指定默认值之后, 是从表定义里面是看不到默认值的,
    #
    id: peewee.AutoField = peewee.AutoField()
    #
    exchange: str = peewee.CharField()
    symbol: str = peewee.CharField()
    interval: str = peewee.CharField()
    period: int = peewee.IntegerField()
    datetime: "datetime.datetime" = peewee.DateTimeField()
    #
    open_price: float = peewee.FloatField()
    high_price: float = peewee.FloatField()
    low_price: float = peewee.FloatField()
    close_price: float = peewee.FloatField()
    volume: float = peewee.FloatField()
    turnover: float = peewee.FloatField()
    open_interest: float = peewee.FloatField()
    settle_price: float = peewee.FloatField()
    #
    tqflag: int = peewee.IntegerField()
    setcode: int = peewee.IntegerField()
    code: str = peewee.CharField()

    class Meta:
        database: playhouse.sqlite_ext.SqliteExtDatabase = _tdx_db
        indexes: tuple = ((("exchange", "symbol", "period", "datetime"), True),)

    def load_and_convert(
        exchange: vnpy.trader.constant.Exchange,
        symbol: str,
        period: TdxPeriod,
        begin_date: "datetime.date",
        last_date: "datetime.date",
    ) -> List[vnpy.trader.object.BarData]:
        """"""
        def reinit_by_date(
            cache: Dict[datetime.date, Dict[datetime.time, datetime.datetime]],
            exchange: vnpy.trader.constant.Exchange,
            symbol: str,
            period: TdxPeriod,
            business_day: datetime.date,
        ) -> None:
            if cache.get(business_day):
                return
            assert period == TdxPeriod._12_5sec
            dttms: Optional[List["datetime.datetime"]] = ChineseMainlandTimeConfig.calc_for_business_day(
                exchange=exchange,
                symbol=symbol,
                unit_second=5 if (period == TdxPeriod._12_5sec) else None,
                business_day=business_day,
            )
            assert dttms is not None
            cache[business_day] = {dttms[i].time(): dttms[i - 1] for i in range(1, len(dttms))}

        cache: Dict[datetime.date, Dict[datetime.time, datetime.datetime]] = {}

        # 从通达信download下来的数据, 其时间戳中, 日期戳是业务日, 时间戳是真实时间段的结束时刻,
        select_count: int = (
            DbBarDataTdx.select()
            .where(
                True
                & (DbBarDataTdx.exchange == exchange.value)
                & (DbBarDataTdx.symbol == symbol)
                & (DbBarDataTdx.interval == period.value)
                & (DbBarDataTdx.datetime >= datetime.datetime.combine(begin_date, datetime.time.min))
                & (DbBarDataTdx.datetime <= datetime.datetime.combine(last_date, datetime.time.max))
            ).count()
        )

        model_select: peewee.ModelSelect = (
            DbBarDataTdx.select()
            .where(
                True
                & (DbBarDataTdx.exchange == exchange.value)
                & (DbBarDataTdx.symbol == symbol)
                & (DbBarDataTdx.interval == period.value)
                & (DbBarDataTdx.datetime >= datetime.datetime.combine(begin_date, datetime.time.min))
                & (DbBarDataTdx.datetime <= datetime.datetime.combine(last_date, datetime.time.max))
            )
            .order_by(
                DbBarDataTdx.datetime
            )
        )

        lines: List[vnpy.trader.object.BarData] = []

        for db_line in model_select:
            db_line: DbBarDataTdx = db_line
            reinit_by_date(
                cache=cache,
                exchange=vnpy.trader.constant.Exchange(value=db_line.exchange),
                symbol=db_line.symbol,
                period=TdxPeriod(value=db_line.interval),
                business_day=db_line.datetime.date(),
            )
            dttm: datetime.datetime = cache[db_line.datetime.date()][db_line.datetime.time()]

            bar: vnpy.trader.object.BarData = vnpy.trader.object.BarData(
                gateway_name="tdx_lib",
                symbol=db_line.symbol,
                exchange=vnpy.trader.constant.Exchange(value=db_line.exchange),
                datetime=vnpy.trader.database.convert_tz(dt=dttm),
                interval=vnpy.trader.constant.Interval.SECOND05,
                volume=db_line.volume,
                turnover=db_line.turnover,
                open_interest=db_line.open_interest,
                open_price=db_line.open_price,
                high_price=db_line.high_price,
                low_price=db_line.low_price,
                close_price=db_line.close_price,
            )
            lines.append(bar)

        lines.sort(key=lambda _: _.datetime)

        assert len({_.datetime for _ in lines}) == len(lines) == select_count

        return lines


class DbBarDataTdxDatabase(object):
    """通过lib将Tdx的kBar数据download到SQLite的数据库接口"""

    def __init__(self) -> None:
        """"""
        self.db: playhouse.sqlite_ext.SqliteExtDatabase = _tdx_db
        self.db.connect()
        self.db.create_tables([DbBarDataTdx])

    def reinitialize(self, sqlite_filename: Optional[str] = None):
        """重新初始化SQLite文件"""
        if sqlite_filename:
            db = DbBarDataTdx._meta.database
            assert isinstance(db, playhouse.sqlite_ext.SqliteExtDatabase)
            assert id(self.db) == id(db)
            #
            path: str = str(vnpy.trader.utility.get_file_path(sqlite_filename))
            self.db.init(database=path)
            #
            self.db.connect()
            self.db.create_tables([DbBarDataTdx])


_tdx_database: DbBarDataTdxDatabase = None


def get_tdx_database() -> DbBarDataTdxDatabase:
    """"""
    # Return database object if already inited
    global _tdx_database
    if _tdx_database:
        return _tdx_database

    _tdx_database = DbBarDataTdxDatabase()
    return _tdx_database


_path: str = str(vnpy.trader.utility.get_file_path("database_backtest.db"))
_db: playhouse.sqlite_ext.SqliteExtDatabase = playhouse.sqlite_ext.SqliteExtDatabase(_path)


class DbCtaBacktestData(vnpy.tools.database_cta.DbCtaBacktestData):
    """"""

    class Meta:
        database: playhouse.sqlite_ext.SqliteExtDatabase = _db


class CtaBacktestDatabase(object):
    """回测用途的SQLite数据库接口"""

    def __init__(self) -> None:
        """"""
        self.db: playhouse.sqlite_ext.SqliteExtDatabase = _db
        self.db.connect()
        self.db.create_tables([DbCtaBacktestData])

    def reinitialize(self, sqlite_filename: Optional[str] = None):
        """重新初始化SQLite文件"""
        if sqlite_filename:
            db = DbCtaBacktestData._meta.database
            assert isinstance(db, playhouse.sqlite_ext.SqliteExtDatabase)
            assert id(self.db) == id(db)
            #
            path: str = str(vnpy.trader.utility.get_file_path(sqlite_filename))
            self.db.init(database=path)
            #
            self.db.connect()
            self.db.create_tables([DbCtaBacktestData])


_database: CtaBacktestDatabase = None


def get_database() -> CtaBacktestDatabase:
    """"""
    # Return database object if already inited
    global _database
    if _database:
        return _database

    _database = CtaBacktestDatabase()
    return _database


def do_tdx_dump(
    filename_src: Optional[str],
    filename_dst: Optional[str],
    exchange: vnpy.trader.constant.Exchange,
    symbol: str,
    period: TdxPeriod,
    begin_date: datetime.date,
    last_date: datetime.date,
) -> None:
    """将5秒的kBar转换并保存到另一个数据库中"""
    db_src: DbBarDataTdxDatabase = get_tdx_database()
    db_dst: vnpy.trader.database.BaseDatabase = vnpy.trader.database.get_database()

    db_src.reinitialize(sqlite_filename=filename_src)
    vnpy_sqlite_reinitialize(db0=db_dst.db, sqlite_filename=filename_dst)

    logging.info(f"db_src: {db_src.db.database}")
    logging.info(f"db_dst: {db_dst.db.database}")

    lines: List[vnpy.trader.object.BarData] = DbBarDataTdx.load_and_convert(
        exchange=exchange,
        symbol=symbol,
        period=period,
        begin_date=begin_date,
        last_date=last_date,
    )

    logging.info(f"len_lines: {len(lines)}")

    db_dst.save_bar_data(bars=lines)


def do_dump(filename_src: str, filename_dst: str):
    """"""
    db_src: vnpy.tools.database_cta.CtaBacktestDatabase = vnpy.tools.database_cta.get_database()
    db_dst: CtaBacktestDatabase = get_database()

    db_src.reinitialize(sqlite_filename=filename_src)
    db_dst.reinitialize(sqlite_filename=filename_dst)

    logging.info(f"db_src: {db_src.db.database}")
    logging.info(f"db_dst: {db_dst.db.database}")

    model_select_src: peewee.ModelSelect = (
        vnpy.tools.database_cta.DbCtaBacktestData.select()
        .where(True)
    )

    count: int = 0

    db_line_batch: List[vnpy.tools.database_cta.DbCtaBacktestData] = []

    for db_line in model_select_src:
        db_line: vnpy.tools.database_cta.DbCtaBacktestData = db_line
        db_line_batch.append(db_line)
        #
        if 50 <= len(db_line_batch):
            ids = [_.id for _ in db_line_batch]
            for db_line in db_line_batch:
                db_line.id = None
            #
            DbCtaBacktestData.bulk_create_on_conflict_replace(model_list=db_line_batch)
            vnpy.tools.database_cta.DbCtaBacktestData.delete().where(vnpy.tools.database_cta.DbCtaBacktestData.id.in_(ids)).execute()
            db_line_batch.clear()
            #
            count += len(db_line_batch)

    if db_line_batch:
        ids = [_.id for _ in db_line_batch]
        for db_line in db_line_batch:
            db_line.id = None
        #
        DbCtaBacktestData.bulk_create_on_conflict_replace(model_list=db_line_batch)
        vnpy.tools.database_cta.DbCtaBacktestData.delete().where(vnpy.tools.database_cta.DbCtaBacktestData.id.in_(ids)).execute()
        db_line_batch.clear()
        #
        count += len(db_line_batch)

    logging.info(f"dump, count={count}")


if __name__ == "__main__bak":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s - %(message)s",)

    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--src", type=str)
    parser.add_argument("-d", "--dst", type=str, default="database_ctabacktest.db")
    args = parser.parse_args()

    filename_src: str = args.src
    filename_dst: str = args.dst

    do_dump(filename_src=filename_src, filename_dst=filename_dst)


if __name__ == "__main__":
    # UPDATE dbbardatatdx SET exchange='CZCE', symbol='MA2405' WHERE exchange='' AND symbol='';
    # SELECT DATE(datetime), COUNT(*) FROM dbbardatatdx GROUP BY DATE(datetime);
    # SELECT DATE(datetime), COUNT(*) FROM dbbardata WHERE interval='multisec' GROUP BY DATE(datetime);
    # SELECT strftime('%Y-%m-%d %H', datetime) AS text_key, COUNT(*) FROM dbbardata WHERE interval='multisec' GROUP BY text_key;
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s - %(message)s",)
    do_tdx_dump(
        None,
        None,
        exchange=vnpy.trader.constant.Exchange.CZCE,
        symbol="MA2405",
        period=TdxPeriod._12_5sec,
        begin_date=datetime.date(year=2024, month=1, day=15),
        last_date=datetime.date(year=2024, month=1, day=24),
    )
